from typing import Dict, Any
from langgraph.graph import StateGraph
from intelli_port.integrations.payments import pay_authorize, pay_capture, pay_refund
from pydantic import BaseModel, Field

class ServicePay(BaseModel):
    type: str = Field(default="pay")
    payment_status: str
    txn_id: str | None = None

class ServiceRefund(BaseModel):
    type: str = Field(default="refund")
    refund_status: str
    refund_id: str | None = None

def build_payment_subgraph():
    g = StateGraph(dict)

    def exec_payment(state: Dict[str, Any]) -> Dict[str, Any]:
        intent = state.get("intent")
        slots = state.get("slots") or {}
        sess = state.get("session") or {}
        risk = state.get("risk_flags") or {}
        payload = {**slots, "user_id": (sess or {}).get("user_id"), "device_id": (sess or {}).get("device_id"), "risk_flags": risk}
        if intent == "pay":
            auth = pay_authorize(payload)
            cap = pay_capture(auth.get("auth_id", ""), payload)
            svc = ServicePay(type="pay", payment_status=cap.get("status", "captured"), txn_id=cap.get("txn_id")).model_dump()
            return {**state, "service": svc}
        if intent == "refund":
            txn_id = (state.get("service") or {}).get("txn_id") or ""
            rf = pay_refund(txn_id, payload)
            svc = ServiceRefund(type="refund", refund_status=rf.get("status", "succeeded"), refund_id=rf.get("refund_id")).model_dump()
            return {**state, "service": svc}
        return {**state, "service": {"type": "pay", "payment_status": "skipped"}}

    g.add_node("payment", exec_payment)
    g.set_entry_point("payment")
    return g.compile()